from timm.scheduler.cosine_lr import CosineLRScheduler